# Kestrel: An XMPP-based Many-Task Computing Scheduler
# Author: Lance Stout <lancestout@gmail.com>
#
# Copyright 2009-2010 Clemson University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import yaml

from kestrel.core.datatypes import Worker, Job, Task
from kestrel.manager import WorkerPool, JobQueue

class Interface(object):
    def __init__(self, kernel):
        """
        The manager interface listens for XMPP messages that
        contain job requests or other similar administrative
        requests.

        Events listened for by this module:
          message
          job/approved
          job/rejected
          job/finished
          task/send
          profile/approved

        Events raised by this module:
          job/received
          job/cancel
          profile/received
          status/requested
          job/task/start
          job/task/end
          job/task/rejected
        """
        self.kernel = kernel
        self.events = {'message': self.message,
                       'job/approved': self.job_approved,
                       'job/rejected': self.job_rejected,
                       'job/finished': self.job_finished,
                       'task/send': self.task_send,
                       'status/send': self.status_send,
                       'profile/approved': self.profile_approved,
                       'task/send/cancel': self.cancel_task}
        self.kernel.register(self)
        self.msg_events = {'job': 'job/received',
                           'cancel': 'job/cancel',
                           'profile': 'profile/received',
                           'status': 'status/requested',
                           'task/started': 'job/task/start',
                           'task/finished': 'job/task/end',
                           'task/rejected': 'job/task/rejected'}
        self.task_events = ['job/task/start', 'job/task/end', 'job/task/rejected']


    def message(self, data=None):
        """
        Raise events if a received message has a known format.

        In this particular implementation, messages whose body
        is a dictionary written in YAML or JSON and
        contains a 'type' entry.
        """
        body = data['body']
        body['source'] = data['from']
        try:
            msg_type = body.get('type')
            event_name = self.msg_events.get(msg_type, None)
            if event_name is not None:
                if event_name in self.task_events:
                    self.kernel.event(event_name, Task(body))
                else:
                    self.kernel.event(event_name, body)
        except:
            pass


    def job_approved(self, job):
        """Send a job/approved message."""
        msg = {'type': 'job/approved',
               'job_id': job.job_id}

        self.kernel.event('message/send', {'to': job.source, 'body': str(msg)})


    def job_rejected(self, job):
        """Send a job/rejected message."""
        msg = {'type': 'job/rejected',
               'job_id': job.job_id}

        self.kernel.event('message/send', {'to': job.source, 'body': str(msg)})


    def task_send(self, data):
        """Send a task to a worker."""
        worker = data['to']
        task = data['task']

        msg = {'type': 'task',
               'job_id': task.job_id,
               'task_id': task.task_id,
               'command': task.command}
        if task.cleanup != '':
            msg['cleanup'] = task.cleanup

        self.kernel.event('message/send', {'to': worker, 'body': str(msg)})


    def job_finished(self, job):
        """Alert the job owner that the job has finished."""
        msg = {'type': 'job/finished',
               'job_id': job.job_id}
        self.kernel.event('message/send', {'to': job.source, 'body': str(msg)})


    def status_send(self, status):
        """Send a status report."""
        msg = {'type':'status',
               status['type']: status['status']}
        self.kernel.event('message/send', {'to': status['to'], 'body': str(msg)})


    def profile_approved(self, worker):
        """Notify the worker that its profile was accepted."""
        msg = {'type':'profile/approved'}
        self.kernel.event('message/send', {'to':worker.id, 'body': str(msg)})


    def cancel_task(self, data):
        """Notify a worker to cancel its task."""
        worker = data['to']
        task = data['task']
        msg = {'type':'cancel',
               'job_id': task.job_id,
               'task_id': task.task_id}
        self.kernel.event('message/send', {'to':worker, 'body':str(msg)})
